import time
from typing import Optional, Callable, Dict, Any

class ProgressReporter:
    """进度报告工具类"""
    
    def __init__(self, total_steps: int, description: str = "", callback: Optional[Callable] = None):
        """
        初始化进度报告器
        
        Args:
            total_steps: 总步骤数
            description: 进度描述
            callback: 进度回调函数
        """
        self.total_steps = max(1, total_steps)  # 避免除零错误
        self.current_step = 0
        self.description = description
        self.callback = callback
        self.start_time = time.time()
        self.last_update_time = self.start_time
        self.update_interval = 0.1  # 更新间隔（秒）
    
    def update(self, step: int = 1, force: bool = False) -> None:
        """
        更新进度
        
        Args:
            step: 步进数量
            force: 是否强制更新（忽略时间间隔限制）
        """
        self.current_step = min(self.current_step + step, self.total_steps)
        current_time = time.time()
        
        # 检查是否需要更新（考虑时间间隔）
        if not force and (current_time - self.last_update_time) < self.update_interval:
            return
        
        if self.callback:
            # 计算进度百分比
            progress = (self.current_step / self.total_steps) * 100
            
            # 计算预计剩余时间
            elapsed_time = current_time - self.start_time
            if self.current_step > 0:
                remaining_steps = self.total_steps - self.current_step
                time_per_step = elapsed_time / self.current_step
                eta = remaining_steps * time_per_step
            else:
                eta = 0
            
            # 构建状态信息
            status = {
                "progress": progress,
                "current_step": self.current_step,
                "total_steps": self.total_steps,
                "description": self.description,
                "elapsed_time": elapsed_time,
                "eta": eta
            }
            
            # 调用回调函数
            self.callback(status)
        
        self.last_update_time = current_time
    
    def done(self) -> None:
        """完成进度报告"""
        if self.callback:
            self.current_step = self.total_steps
            self.update(force=True)

class NodeProgressReporter(ProgressReporter):
    """ComfyUI节点进度报告器"""
    
    def __init__(self, total_steps: int, description: str = "", node = None):
        """
        初始化节点进度报告器
        
        Args:
            total_steps: 总步骤数
            description: 进度描述
            node: ComfyUI节点实例
        """
        super().__init__(total_steps, description, self._node_callback)
        self.node = node
    
    def _node_callback(self, status: Dict[str, Any]) -> None:
        """节点进度回调"""
        if not self.node:
            return
        
        # 格式化进度信息
        progress = f"{status['description']} - {status['progress']:.1f}%"
        if status['eta'] > 0:
            progress += f" (剩余时间: {status['eta']:.1f}秒)"
        
        # 更新节点状态
        if hasattr(self.node, 'update_status'):
            self.node.update_status({"value": progress}) 